package com.yeming.flink.mudlog.isolationforest

import org.apache.flink.api.scala.{DataSet, _}

import scala.util.Random

sealed trait ITree
case class ITreeBranch(left: ITree, right: ITree, split_column: Int, split_value: Double) extends ITree
case class ITreeLeaf(size: Long) extends ITree

/**
 * 孤立森林样例类
 * @param num_samples
 * @param trees
 */
case class IsolationForest(num_samples: Long, trees: Array[ITree]) {

  def predict(x:Array[Double]): Double = {
    val predictions = trees.map(s=>pathLength(x, s, 0)).toList
    math.pow(2, -(predictions.sum/predictions.size)/cost(num_samples))
  }

  def cost(num_items:Long): Int =
    (2*(math.log(num_items-1)+0.5772156649)-(2*(num_items-1)/num_items)).toInt

  @scala.annotation.tailrec
  final def pathLength(x:Array[Double], tree:ITree, path_length:Int): Double ={
    tree match{
      case ITreeLeaf(size) =>
        if (size > 1) path_length + cost(size)
        else path_length + 1

      case ITreeBranch(left, right, split_column, split_value) =>
        val sample_value = x(split_column)

        if (sample_value < split_value)
          pathLength(x, left, path_length+1)
        else
          pathLength(x, right, path_length+1)
    }
  }
}
/**
 * 孤立森林
 */
object IsolationForest {


  def buildForest(data: DataSet[Array[Double]], numTrees: Int = 2, subSampleSize: Int = 256, seed: Long = Random.nextLong) : IsolationForest = {

    println("开始构建孤立森林……")
    val numSamples = data.count()
    val numColumns = data.first(1).collect()(0).length
    val maxHeight = math.ceil(math.log(subSampleSize)).toInt
    val trees = Array.fill[ITree](numTrees)(ITreeLeaf(1))

    val trainedTrees = trees.map(s=>growTree(getRandomSubsample(data, subSampleSize/numSamples.toDouble, seed), maxHeight, numColumns))

    IsolationForest(numSamples, trainedTrees)
  }



  def growTree(data: DataSet[Array[Double]], maxHeight:Int, numColumns:Int, currentHeight:Int = 0): ITree = {
    println("开始构建孤立树……")
    println(s"当前树的最大高度为：${maxHeight}, 维度：${numColumns}")
    val numSamples = data.count()
    //递归终止条件，当前树高大于maxHeight或数据量不大于1
    if(currentHeight>=maxHeight || numSamples <= 1){
      return  ITreeLeaf(numSamples)
    }
    //随机选择特征列
    val split_column = Random.nextInt(numColumns)
    val column = data.map(s => (s(split_column),1))
    //随机选择该特征列中的值split_value，用于分割样本
    val col_min = column.min(0).collect()(0)._1
    val col_max = column.max(0).collect()(0)._1
    val split_value = col_min + Random.nextDouble() * (col_max - col_min)
    //小于分割值的成为左子树，反之右子树
    val X_left = data.filter(s => s(split_column) < split_value)
    val X_right = data.filter(s => s(split_column) >= split_value)

    //递归
    new ITreeBranch(growTree(X_left, maxHeight, numColumns, currentHeight + 1),
      growTree(X_right, maxHeight, numColumns, currentHeight + 1), split_column, split_value)
  }

  def getRandomSubsample(data: DataSet[Array[Double]], sampleRatio: Double, seed: Long = Random.nextLong): DataSet[Array[Double]]  = {
    val rd = new Random(seed)
    data.filter(data => rd.nextDouble() < sampleRatio)
  }
}
